import pandas as pd
import numpy as np
from sklearn.model_selection import train_test_split, cross_val_score, KFold
from sklearn.preprocessing import StandardScaler
from sklearn.ensemble import RandomForestRegressor, GradientBoostingRegressor
from sklearn.linear_model import LogisticRegression
from sklearn.impute import SimpleImputer
import xgboost as xgb
from sklearn.metrics import mean_squared_error, r2_score, mean_absolute_error
import logging
import joblib
import matplotlib.pyplot as plt
from datetime import datetime
from pathlib import Path


class OlympicPredictor:
    def __init__(self):
        self.logger = self._setup_logger()
        self.scaler = StandardScaler()
        self.cv = KFold(n_splits=5, shuffle=True, random_state=42)
        
    def _setup_logger(self):
        logger = logging.getLogger('OlympicPredictor')
        logger.setLevel(logging.INFO)
        handler = logging.StreamHandler()
        formatter = logging.Formatter('%(asctime)s - %(name)s - %(levelname)s - %(message)s')
        handler.setFormatter(formatter)
        logger.addHandler(handler)
        return logger

    def load_data(self):
        """加载并预处理数据"""
        try:
            # 加载基础数据
            self.athlete_data = pd.read_csv('features/athlete_features.csv')
            self.country_data = pd.read_csv('features/country_features.csv')
            self.medal_encoding = pd.read_csv('features/Medal_encoding.csv')
            self.noc_encoding = pd.read_csv('features/NOC_encoding.csv')
            
            # 加载未获奖国家特征
            self.non_medal_features = pd.read_csv('features/non_medal_countries_features.csv')
            
            # 只保留2020和2024年的数据
            recent_years = [2020, 2024]
            self.athlete_data = self.athlete_data[
                self.athlete_data['last_year'].isin(recent_years)
            ]
            
            # 更新country_data的统计数据
            self._update_country_stats()
            
            # 添加数据验证
            print("数据加载验证:")
            print(f"country_data shape: {self.country_data.shape}")
            print(f"列名: {self.country_data.columns.tolist()}")

            
            self.logger.info("数据加载成功")
            return True
        except Exception as e:
            self.logger.error(f"数据加载失败: {str(e)}")
            return False

    def _update_country_stats(self):
        """更新国家统计数据，基于最近的比赛"""
        stats = {}
        
        for country in self.country_data['country'].unique():
            country_info = self.country_data[
                self.country_data['country'] == country
            ].iloc[0]
            
            # 计算最近的表现
            recent_stats = {
                'country': country,
                'recent_gold': country_info['total_gold'] if country_info['last_year'] >= 2020 else 0,
                'recent_total': country_info['total_medals'] if country_info['last_year'] >= 2020 else 0,
                'athlete_count': len(self.athlete_data[self.athlete_data['country'] == country]),
                'avg_performance': self.athlete_data[
                    self.athlete_data['country'] == country
                ]['avg_performance'].mean(),
                'sport_diversity': self.athlete_data[
                    self.athlete_data['country'] == country
                ]['sport_count'].mean(),
                'is_host_2028': 1 if country == 'United States' else 0  # 2028年主办国
            }
            
            stats[country] = recent_stats
            
        self.country_stats = pd.DataFrame.from_dict(stats, orient='index')

    def prepare_features(self):
        """准备预测特征"""
        features = pd.DataFrame()
        
        # 基础特征
        features['athlete_count'] = self.country_stats['athlete_count']
        features['avg_performance'] = self.country_stats['avg_performance'].fillna(0)
        features['sport_diversity'] = self.country_stats['sport_diversity'].fillna(1)
        
        # 历史表现特征 - 修改这部分，分开处理金牌和总奖牌
        features['recent_gold'] = self.country_stats['recent_gold']
        features['recent_total'] = self.country_stats['recent_total']
        features['gold_ratio'] = (
            self.country_stats['recent_gold'] / 
            self.country_stats['recent_total'].clip(lower=1)
        )
        
        # 主办国效应
        features['is_host'] = self.country_stats['is_host_2028']
        
        # 标准化特征
        features_scaled = self.scaler.fit_transform(features)
        features_scaled = pd.DataFrame(
            features_scaled, 
            columns=features.columns,
            index=features.index
        )
        
        return features_scaled

    def evaluate_model_performance(self, features, gold_targets, total_targets):
        """评估模型性能"""
        # 转换特征为numpy数组，避免特征名称警告
        X = features.values
        
        # 交叉验证评估
        gold_cv_scores = cross_val_score(
            self.gold_model, X, gold_targets, 
            cv=self.cv, scoring='neg_mean_squared_error'
        )
        total_cv_scores = cross_val_score(
            self.total_model, X, total_targets, 
            cv=self.cv, scoring='neg_mean_squared_error'
        )
        
        # 计算RMSE
        gold_rmse = np.sqrt(-gold_cv_scores.mean())
        total_rmse = np.sqrt(-total_cv_scores.mean())
        
        # R2分数
        gold_r2 = self.gold_model.score(X, gold_targets)
        total_r2 = self.total_model.score(X, total_targets)
        
        # 计算MAE
        gold_predictions = self.gold_model.predict(X)
        total_predictions = self.total_model.predict(X)
        gold_mae = mean_absolute_error(gold_targets, gold_predictions)
        total_mae = mean_absolute_error(total_targets, total_predictions)
        
        return {
            'gold_medals': {
                'rmse': gold_rmse,
                'r2': gold_r2,
                'mae': gold_mae,
                'cv_scores': gold_cv_scores
            },
            'total_medals': {
                'rmse': total_rmse,
                'r2': total_r2,
                'mae': total_mae,
                'cv_scores': total_cv_scores
            }
        }
    
    def assess_prediction_uncertainty(self, features):
        """评估预测的不确定性"""
        X = features.values
        
        # 使用随机森林的所有树进行预测
        gold_predictions = np.array([
            tree.predict(X) for tree in self.gold_model.estimators_
        ])
        total_predictions = np.array([
            tree.predict(X) for tree in self.total_model.estimators_
        ])
        
        # 计算预测的标准差（作为不确定性的度量）
        gold_std = np.std(gold_predictions, axis=0)
        total_std = np.std(total_predictions, axis=0)
        
        # 计算变异系数（CV）
        gold_mean = np.mean(gold_predictions, axis=0)
        total_mean = np.mean(total_predictions, axis=0)
        gold_cv = np.divide(gold_std, gold_mean, where=gold_mean!=0)
        total_cv = np.divide(total_std, total_mean, where=total_mean!=0)
        
        return pd.DataFrame({
            'country': features.index,
            'gold_std': gold_std,
            'gold_cv': gold_cv,
            'total_std': total_std,
            'total_cv': total_cv
        })
    
    def evaluate_prediction_intervals(self, features, gold_targets, total_targets):
        """评估预测区间的可靠性"""
        X = features.values
        
        # 生成预测区间
        gold_predictions_all = np.array([
            tree.predict(X) for tree in self.gold_model.estimators_
        ])
        total_predictions_all = np.array([
            tree.predict(X) for tree in self.total_model.estimators_
        ])
        
        # 计算90%预测区间
        gold_lower = np.percentile(gold_predictions_all, 5, axis=0)
        gold_upper = np.percentile(gold_predictions_all, 95, axis=0)
        total_lower = np.percentile(total_predictions_all, 5, axis=0)
        total_upper = np.percentile(total_predictions_all, 95, axis=0)
        
        # 计算实际值落在预测区间内的比例
        gold_coverage = np.mean(
            (gold_targets >= gold_lower) & (gold_targets <= gold_upper)
        )
        total_coverage = np.mean(
            (total_targets >= total_lower) & (total_targets <= total_upper)
        )
        
        return {
            'gold_coverage': gold_coverage,
            'total_coverage': total_coverage
        }
    
    def analyze_feature_stability(self, features, n_iterations=100):
        """分析特征重要性的稳定性"""
        X = features.values
        gold_importance_matrix = np.zeros((n_iterations, features.shape[1]))
        total_importance_matrix = np.zeros((n_iterations, features.shape[1]))
        
        for i in range(n_iterations):
            # 随机抽样
            sample_idx = np.random.choice(
                X.shape[0], 
                size=int(X.shape[0] * 0.8), 
                replace=True
            )
            X_sample = X[sample_idx]
            gold_targets_sample = self.country_stats['recent_gold'].iloc[sample_idx]
            total_targets_sample = self.country_stats['recent_total'].iloc[sample_idx]
            
            # 训练模型
            gold_model = RandomForestRegressor(
                n_estimators=200, random_state=i
            ).fit(X_sample, gold_targets_sample)
            total_model = RandomForestRegressor(
                n_estimators=200, random_state=i
            ).fit(X_sample, total_targets_sample)
            
            # 记录特征重要性
            gold_importance_matrix[i] = gold_model.feature_importances_
            total_importance_matrix[i] = total_model.feature_importances_
        
        # 计算特征重要性的统计量
        feature_stability = pd.DataFrame({
            'feature': features.columns,
            'gold_importance_mean': np.mean(gold_importance_matrix, axis=0),
            'gold_importance_std': np.std(gold_importance_matrix, axis=0),
            'total_importance_mean': np.mean(total_importance_matrix, axis=0),
            'total_importance_std': np.std(total_importance_matrix, axis=0)
        })
        
        return feature_stability

    def train_models(self, features, gold_targets, total_targets):
        """训练预测模型并评估性能"""
        # 训练模型
        self.gold_model = RandomForestRegressor(
            n_estimators=200,
            max_depth=15,
            min_samples_split=5,
            random_state=42
        )
        
        self.total_model = RandomForestRegressor(
            n_estimators=200,
            max_depth=15,
            min_samples_split=5,
            random_state=42
        )
        
        # 转换特征为numpy数组
        X = features.values
        
        self.gold_model.fit(X, gold_targets)
        self.total_model.fit(X, total_targets)
        
        # 评估模型性能
        performance = self.evaluate_model_performance(features, gold_targets, total_targets)
        uncertainty = self.assess_prediction_uncertainty(features)
        interval_reliability = self.evaluate_prediction_intervals(
            features, gold_targets, total_targets
        )
        feature_stability = self.analyze_feature_stability(features)
        
        return {
            'performance': performance,
            'uncertainty': uncertainty,
            'interval_reliability': interval_reliability,
            'feature_stability': feature_stability
        }

    def predict_2028_olympics(self):
        """预测2028年奥运会成绩，包括预测区间"""
        try:
            # 准备特征
            features = self.prepare_features()
            
            # 预测
            gold_predictions = self.gold_model.predict(features)
            total_predictions = self.total_model.predict(features)
            
            # 计算预测区间
            gold_predictions_all = np.array([tree.predict(features) 
                                           for tree in self.gold_model.estimators_])
            total_predictions_all = np.array([tree.predict(features) 
                                            for tree in self.total_model.estimators_])
            
            # 使用更窄的百分位数范围来计算预测区间
            gold_lower = np.percentile(gold_predictions_all, 40, axis=0)  # 从25改为40
            gold_upper = np.percentile(gold_predictions_all, 60, axis=0)  # 从75改为60
            total_lower = np.percentile(total_predictions_all, 40, axis=0)  # 从25改为40
            total_upper = np.percentile(total_predictions_all, 60, axis=0)  # 从75改为60
            
            # 整理预测结果
            predictions = {}
            for idx, country in enumerate(features.index):
                # 使用历史表现来约束预测区间
                country_data = self.country_data[self.country_data['country'] == country]
                recent_total = country_data['total_medals'].iloc[-1] if len(country_data) > 0 else 0
                recent_gold = country_data['total_gold'].iloc[-1] if len(country_data) > 0 else 0
                
                # 设定更严格的波动范围（从±30%改为±20%）
                total_range = max(recent_total * 0.2, 5)  # 从10改为5
                gold_range = max(recent_gold * 0.2, 3)    # 从5改为3
                
                # 计算基础预测值
                base_total = round(total_predictions[idx])
                base_gold = round(gold_predictions[idx])
                
                predictions[country] = {
                    'predicted_total': max(0, base_total),
                    'predicted_gold': max(0, base_gold),
                    'total_lower': max(0, round(base_total - total_range)),
                    'total_upper': round(base_total + total_range),
                    'gold_lower': max(0, round(base_gold - gold_range)),
                    'gold_upper': round(base_gold + gold_range),
                    'athlete_count': self.country_stats.loc[country, 'athlete_count']
                }
                
                # 对于历史成绩稳定的国家，使用更窄的区间
                if len(country_data) > 0 and country_data['avg_medals'].iloc[0] > 0:
                    medals_std = country_data['total_medals'].std()
                    if medals_std / country_data['avg_medals'].iloc[0] < 0.3:  # 历史表现稳定
                        predictions[country]['total_lower'] = max(0, round(base_total - total_range * 0.7))
                        predictions[country]['total_upper'] = round(base_total + total_range * 0.7)
                        predictions[country]['gold_lower'] = max(0, round(base_gold - gold_range * 0.7))
                        predictions[country]['gold_upper'] = round(base_gold + gold_range * 0.7)
                
                # 确保金牌预测区间不超过总奖牌预测区间
                predictions[country]['gold_upper'] = min(
                    predictions[country]['gold_upper'],
                    predictions[country]['total_upper']
                )
                predictions[country]['gold_lower'] = min(
                    predictions[country]['gold_lower'],
                    predictions[country]['total_lower']
                )
            
            # 调整预测结果确保合理性
            predictions = self._adjust_predictions_with_intervals(predictions)
            
            # 将预测结果转换为DataFrame以便可视化
            results_df = pd.DataFrame.from_dict(predictions, orient='index')
            
            # 1. 创建前30名国家的奖牌预测可视化
            top_30 = results_df.nlargest(30, 'predicted_total')
            
            plt.figure(figsize=(15, 10))
            x = np.arange(len(top_30))
            width = 0.35
            
            # 绘制总奖牌数和金牌数的柱状图
            plt.bar(x - width/2, top_30['predicted_total'], width, 
                    label='Predicted Total Medals', color='blue', alpha=0.7)
            plt.bar(x + width/2, top_30['predicted_gold'], width,
                    label='Predicted Gold Medals', color='gold', alpha=0.7)
            
            # 添加预测区间
            plt.errorbar(x - width/2, top_30['predicted_total'],
                        yerr=[(top_30['predicted_total'] - top_30['total_lower']),
                              (top_30['total_upper'] - top_30['predicted_total'])],
                        fmt='none', color='red', alpha=0.5)
            plt.errorbar(x + width/2, top_30['predicted_gold'],
                        yerr=[(top_30['predicted_gold'] - top_30['gold_lower']),
                              (top_30['gold_upper'] - top_30['predicted_gold'])],
                        fmt='none', color='darkred', alpha=0.5)
            
            plt.title('2028 Olympics Medal Predictions - Top 30 Countries')
            plt.xlabel('Country')
            plt.ylabel('Number of Medals')
            plt.xticks(x, top_30.index, rotation=45, ha='right')
            plt.legend()
            plt.grid(True, alpha=0.3)
            
            # Save the plot
            plt.tight_layout()
            plt.savefig('2028_medals_prediction.png', dpi=300, bbox_inches='tight')
            plt.close()
            
            # 2. 保存预测结果到CSV文件，按总奖牌数排序
            results_df = results_df.sort_values('predicted_total', ascending=False)
            results_df.to_csv('2028_olympics_predictions.csv')
            
            return predictions
        
        except Exception as e:
            self.logger.error(f"预测失败: {str(e)}")
            raise

    def _adjust_predictions_with_intervals(self, predictions):
        """调整预测结果确保合理性，包括预测区间"""
        total_events_2028 = 329  # 2024年的项目数
        total_medals_2028 = 1039  # 使用2024年的实际奖牌数
        total_gold_2028 = total_events_2028  # 金牌数等于项目数
        
        # 计算总预测奖牌数和金牌数
        total_predicted_medals = sum(p['predicted_total'] for p in predictions.values())
        total_predicted_gold = sum(p['predicted_gold'] for p in predictions.values())
        
        # 分别调整总奖牌和金牌数
        if total_predicted_medals != total_medals_2028:
            medals_scale_factor = total_medals_2028 / max(total_predicted_medals, 1)
            for country in predictions:
                predictions[country]['predicted_total'] = round(
                    predictions[country]['predicted_total'] * medals_scale_factor
                )
                # 同时调整预测区间
                predictions[country]['total_lower'] = round(
                    predictions[country]['total_lower'] * medals_scale_factor
                )
                predictions[country]['total_upper'] = round(
                    predictions[country]['total_upper'] * medals_scale_factor
                )
        
        if total_predicted_gold != total_gold_2028:
            gold_scale_factor = total_gold_2028 / max(total_predicted_gold, 1)
            for country in predictions:
                # 调整金牌数，确保不超过该国总奖牌数
                predictions[country]['predicted_gold'] = min(
                    round(predictions[country]['predicted_gold'] * gold_scale_factor),
                    predictions[country]['predicted_total']
                )
                # 同时调整金牌预测区间
                predictions[country]['gold_lower'] = min(
                    round(predictions[country]['gold_lower'] * gold_scale_factor),
                    predictions[country]['total_lower']
                )
                predictions[country]['gold_upper'] = min(
                    round(predictions[country]['gold_upper'] * gold_scale_factor),
                    predictions[country]['total_upper']
                )
        
        return predictions

    def get_feature_importance(self):
        """获取特征重要性"""
        feature_names = self.prepare_features().columns
        
        gold_importance = pd.DataFrame({
            'feature': feature_names,
            'importance': self.gold_model.feature_importances_
        }).sort_values('importance', ascending=False)
        
        total_importance = pd.DataFrame({
            'feature': feature_names,
            'importance': self.total_model.feature_importances_
        }).sort_values('importance', ascending=False)
        
        return {
            'gold_medals': gold_importance,
            'total_medals': total_importance
        }

    def predict_2028_athletes(self):
        """预测2028年各国参赛运动员数量"""
        athlete_predictions = {}
        current_year = 2024  # 当前奥运年
        target_year = 2028  # 目标预测年
        
        for country in self.country_stats.index:
            current_athletes = self.athlete_data[
                self.athlete_data['country'] == country
            ]
            
            if len(current_athletes) == 0:
                continue
            
            # 根据最后参赛年份来估计运动员是否会继续参加2028奥运会
            last_participation = current_athletes['last_year']
            years_to_2028 = target_year - last_participation
            
            # 假设运动员参加超过2届奥运会的概率降低
            likely_to_continue = (years_to_2028 <= 8).sum()  # 最多参加3届奥运会
            
            # 估计新增运动员数量（基于历史趋势）
            recent_athletes = len(current_athletes[
                current_athletes['last_year'] >= 2020
            ])
            historical_growth = recent_athletes * 0.3  # 假设每届增加30%新运动员
            
            # 预测2028年运动员数量
            predicted_athletes = round(likely_to_continue * 0.7 + historical_growth)  # 考虑一定的退役率
            
            # 确保预测数量合理
            min_athletes = max(5, round(recent_athletes * 0.5))  # 最少保持50%的规模
            max_athletes = round(recent_athletes * 1.5)  # 最多增长到1.5倍
            predicted_athletes = max(min_athletes, min(predicted_athletes, max_athletes))
            
            athlete_predictions[country] = predicted_athletes
        
        return athlete_predictions

    def prepare_country_features(self):
        """准备国家特征"""
        # 预测2028年运动员数量
        athlete_predictions_2028 = self.predict_2028_athletes()
        
        for country in self.country_stats.index:
            self.country_stats.loc[country, 'athlete_count'] = athlete_predictions_2028.get(
                country, 
                self.country_stats.loc[country, 'athlete_count']
            )

    def predict_first_medal_probability(self):
        """预测首次获奖可能性"""
        try:
            # 读取从未获得过奖牌的国家列表
            non_medal_countries = pd.read_csv('non_medal_countries.csv')
            predictions = {}
            
            # 获取运动员数据
            for country in non_medal_countries['Country']:
                # 获取该国运动员数据
                country_athletes = self.athlete_data[
                    self.athlete_data['country'] == country
                ]
                
                # 如果国家没有运动员数据，设置默认值
                if len(country_athletes) == 0:
                    predictions[country] = {
                        'probability': 0.0,  # 获奖概率为0
                        'confidence': 'Low',  # 置信度为低
                        'athlete_count': 0,   # 运动员数量为0
                        'unique_sports': 0,   # 运动项目数为0
                        'key_factors': ['从未参加奥运会']  # 特殊标记
                    }
                    continue
                    
                # 计算基础特征
                athlete_count = len(country_athletes)
                unique_sports = country_athletes['sport_count'].mean()
                
                # 计算获奖概率
                probability = self._simple_probability_calculation(
                    athlete_count, 
                    unique_sports
                )
                
                # 确定预测置信度
                confidence = self._simple_confidence_calculation(
                    athlete_count,
                    unique_sports
                )
                
                # 识别关键因素
                key_factors = []
                if athlete_count > 20:
                    key_factors.append("运动员数量充足")
                if unique_sports > 2:
                    key_factors.append("运动项目多样化")
                if not key_factors:  # 如果没有积极因素
                    key_factors.append("需要增加参与度")
                
                predictions[country] = {
                    'probability': round(probability * 100, 2),
                    'confidence': confidence,
                    'athlete_count': athlete_count,
                    'unique_sports': round(unique_sports, 2),
                    'key_factors': key_factors
                }
            
            # 生成预测结果可视化
            self._plot_first_medal_predictions(predictions)
            
            # 保存预测结果
            self.save_first_medal_predictions(predictions)
            
            return predictions
            
        except Exception as e:
            self.logger.error(f"首次获奖预测失败: {str(e)}")
            return None

    def _simple_probability_calculation(self, athlete_count, unique_sports):
        """简化的概率计算"""
        base_probability = 0.0
        
        # 更乐观的运动员数量权重
        if athlete_count > 25:
            base_probability += 0.6  # 增加基础概率
        elif athlete_count > 15:
            base_probability += 0.5
        elif athlete_count > 10:
            base_probability += 0.4
        elif athlete_count > 5:
            base_probability += 0.3
        
        # 更宽松的项目多样性权重
        if unique_sports > 1.2:
            base_probability += 0.3
        elif unique_sports > 1.0:
            base_probability += 0.2
        
        # 额外的发展潜力加成
        if athlete_count > 8:  # 降低发展潜力门槛
            base_probability += 0.1
        
        return min(max(base_probability, 0.0), 1.0)

    def _simple_confidence_calculation(self, athlete_count, unique_sports):
        """简化的置信度计算"""
        # 更乐观的置信度标准
        if athlete_count > 12 and unique_sports > 1.0:  # 降低High门槛
            return "High"
        elif athlete_count > 6 and unique_sports > 1.0:  # 降低Medium门槛
            return "Medium"
        else:
            return "Low"

    def save_first_medal_predictions(self, predictions):
        """保存首次获奖预测结果"""
        try:
            # 创建predictions目录
            Path('predictions').mkdir(exist_ok=True)
            
            # 转换为DataFrame
            df = pd.DataFrame.from_dict(predictions, orient='index')
            
            # 保存到CSV
            df.to_csv('predictions/first_medal_predictions.csv')
            self.logger.info("首次获奖预测结果已保存")
            
        except Exception as e:
            self.logger.error(f"保存预测结果失败: {str(e)}")

    def _plot_first_medal_predictions(self, predictions):
        """生成首次获奖预测可视化"""
        # 设置英文字体
        plt.rcParams['font.sans-serif'] = ['Arial']
        plt.rcParams['axes.unicode_minus'] = False
        
        # 获取数据
        countries = list(predictions.keys())
        probabilities = [pred['probability'] for pred in predictions.values()]
        confidences = [pred['confidence'] for pred in predictions.values()]
        
        # 创建图形
        plt.figure(figsize=(20, 10))
        
        # 设置颜色映射
        color_map = {
            'High': 'red',
            'Medium': 'yellow',
            'Low': 'blue'
        }
        
        # 创建柱状图，根据置信度设置颜色
        bars = plt.bar(countries, probabilities)
        for bar, confidence in zip(bars, confidences):
            bar.set_color(color_map[confidence])
        
        # 设置标签和标题
        plt.xticks(rotation=45, ha='right', fontsize=8)
        plt.yticks(fontsize=10)
        plt.title('First Medal Probability Prediction 2028', fontsize=14, pad=20)
        plt.xlabel('Country', fontsize=12, labelpad=10)
        plt.ylabel('Probability (%)', fontsize=12, labelpad=10)
        
        # 添加图例
        legend_elements = [plt.Rectangle((0,0),1,1, facecolor=color) 
                          for color in color_map.values()]
        plt.legend(legend_elements, color_map.keys(), 
                  title='Confidence Level', loc='upper right')
        
        # 调整布局
        plt.tight_layout()
        
        # 保存图片
        plt.savefig('predictions/first_medal_predictions.png', 
                    dpi=300, 
                    bbox_inches='tight',
                    pad_inches=0.2)
        plt.close()


# 使用示例
if __name__ == "__main__":
    predictor = OlympicPredictor()

    try:
        if predictor.load_data():
            print("\n开始模型训练和评估...")
            
            # 准备特征和目标变量
            features = predictor.prepare_features()
            gold_targets = predictor.country_stats['recent_gold']
            total_targets = predictor.country_stats['recent_total']
            
            # 训练模型并获取评估结果
            evaluation_results = predictor.train_models(
                features, gold_targets, total_targets
            )
            
            # 输出详细的评估结果
            print("\n=== 模型评估结果 ===")
            
            print("\n1. 预测性能:")
            print(f"金牌预测 RMSE: {evaluation_results['performance']['gold_medals']['rmse']:.2f}")
            print(f"金牌预测 R2: {evaluation_results['performance']['gold_medals']['r2']:.2f}")
            print(f"金牌预测 MAE: {evaluation_results['performance']['gold_medals']['mae']:.2f}")
            print(f"\n总奖牌预测 RMSE: {evaluation_results['performance']['total_medals']['rmse']:.2f}")
            print(f"总奖牌预测 R2: {evaluation_results['performance']['total_medals']['r2']:.2f}")
            print(f"总奖牌预测 MAE: {evaluation_results['performance']['total_medals']['mae']:.2f}")
            
            print("\n2. 预测区间可靠性:")
            print(f"金牌预测区间覆盖率: {evaluation_results['interval_reliability']['gold_coverage']:.2%}")
            print(f"总奖牌预测区间覆盖率: {evaluation_results['interval_reliability']['total_coverage']:.2%}")
            
            print("\n3. 特征重要性分析:")
            feature_importance = evaluation_results['feature_stability'].sort_values(
                'gold_importance_mean', 
                ascending=False
            )
            print("\n金牌预测的特征重要性:")
            for idx, row in feature_importance.iterrows():
                print(f"{row['feature']}: {row['gold_importance_mean']:.4f} (±{row['gold_importance_std']:.4f})")
            
            print("\n总奖牌预测的特征重要性:")
            feature_importance = feature_importance.sort_values('total_importance_mean', ascending=False)
            for idx, row in feature_importance.iterrows():
                print(f"{row['feature']}: {row['total_importance_mean']:.4f} (±{row['total_importance_std']:.4f})")
            
            print("\n4. 预测不确定性:")
            uncertainty_stats = evaluation_results['uncertainty'].describe()
            print("\n金牌预测标准差统计:")
            print(uncertainty_stats['gold_std'])
            print("\n总奖牌预测标准差统计:")
            print(uncertainty_stats['total_std'])
            
            # 保存评估结果到文件
            evaluation_path = 'predictions/model_evaluation_results.txt'
            Path('predictions').mkdir(exist_ok=True)
            
            with open(evaluation_path, 'w', encoding='utf-8') as f:
                f.write("=== 奥运会奖牌预测模型评估报告 ===\n\n")
                f.write("1. 预测性能:\n")
                f.write(f"金牌预测 RMSE: {evaluation_results['performance']['gold_medals']['rmse']:.2f}\n")
                f.write(f"金牌预测 R2: {evaluation_results['performance']['gold_medals']['r2']:.2f}\n")
                f.write(f"金牌预测 MAE: {evaluation_results['performance']['gold_medals']['mae']:.2f}\n")
                f.write(f"\n总奖牌预测 RMSE: {evaluation_results['performance']['total_medals']['rmse']:.2f}\n")
                f.write(f"总奖牌预测 R2: {evaluation_results['performance']['total_medals']['r2']:.2f}\n")
                f.write(f"总奖牌预测 MAE: {evaluation_results['performance']['total_medals']['mae']:.2f}\n")

                f.write("\n2. 预测区间可靠性:\n")
                f.write(f"金牌预测区间覆盖率: {evaluation_results['interval_reliability']['gold_coverage']:.2%}\n")
                f.write(f"总奖牌预测区间覆盖率: {evaluation_results['interval_reliability']['total_coverage']:.2%}\n")
                
                f.write("\n3. 特征重要性分析:\n")
                f.write("\n金牌预测的特征重要性:\n")
                for idx, row in feature_importance.iterrows():
                    f.write(f"{row['feature']}: {row['gold_importance_mean']:.4f} (±{row['gold_importance_std']:.4f})\n")
            
            print(f"\n详细的评估结果已保存到: {evaluation_path}")
            
            # 继续原有的预测流程...
            
            # 预测2028年运动员数量
            athlete_predictions = predictor.predict_2028_athletes()
            print("\n2028年预计参赛运动员数量:")
            for country, count in sorted(
                athlete_predictions.items(), 
                key=lambda x: x[1], 
                reverse=True
            )[:20]:
                print(f"{country}: {count}人")
            
            # 预测2028年奥运会结果
            predictions_2028 = predictor.predict_2028_olympics()

            # 输出预测结果
            print("\n2028洛杉矶奥运会预测结果:")
            
            # 按预测总奖牌数排序
            sorted_predictions = dict(
                sorted(
                    predictions_2028.items(),
                    key=lambda x: x[1]['predicted_total'],
                    reverse=True
                )
            )
            
            # 输出前20个国家的预测结果
            print("\n前20个国家的预测结果:")
            for i, (country, pred) in enumerate(sorted_predictions.items(), 1):
                if i <= 20:
                    print(f"\n{i}. {country}:")
                    print(f"   预计总奖牌数: {pred['predicted_total']} (区间: {pred['total_lower']}-{pred['total_upper']})")
                    print(f"   预计金牌数: {pred['predicted_gold']} (区间: {pred['gold_lower']}-{pred['gold_upper']})")
                    print(f"   运动员数量: {pred['athlete_count']}")
                    if country == 'United States':
                        print("   (东道主)")

            # 保存总体预测结果
            results_df = pd.DataFrame.from_dict(predictions_2028, orient='index')
            results_df.to_csv('2028_olympics_predictions.csv')
            print(f"\n完整预测结果已保存到 2028_olympics_predictions.csv")

            # 预测首次获奖可能性
            print("\n开始预测首次获奖可能性...")
            first_medal_predictions = predictor.predict_first_medal_probability()
            
            if first_medal_predictions:
                # 打印首次获奖预测结果
                print("\n首次获奖可能性预测 (前10个国家):")
                predictions_df = pd.DataFrame.from_dict(first_medal_predictions, orient='index')
                predictions_df = predictions_df.sort_values('probability', ascending=False)
                
                for country in predictions_df.head(10).index:
                    pred = first_medal_predictions[country]
                    print(f"\n{country}:")
                    print(f"   获奖概率: {pred['probability']}%")
                    print(f"   预测置信度: {pred['confidence']}")
                    print(f"   运动员数量: {pred['athlete_count']}")
                    print(f"   运动项目数: {pred['unique_sports']}")
                    print(f"   关键因素: {', '.join(pred['key_factors'])}")
            else:
                print("\n没有找到符合条件的首次获奖候选国家")

    except Exception as e:
        print(f"程序执行失败: {str(e)}")
        raise